#include <stdio.h>
#include "../../tdp/tdp_jpegsupport.hpp"
#include "../../tdp/tdp_guard.hpp"
#include "../sb4image.h"
#include "../sb4videoenc.h"

#define RoQ_QUAD_JPEG		0x1012
#define RoQ_QUAD_DJPEG		0x1015

class myVectorWriter : public TDP::terJPEGWriter
{
private:
	terVector<unsigned char> *_v;
public:
	myVectorWriter(terVector<unsigned char> *v)
	{
		_v = v;
	}

	size_t Write(const void *bytes, size_t sz)
	{
		unsigned int offs;
		terVector<unsigned char> oldBytes;
		unsigned char *b;

		offs = _v->Count();
		_v->Resize(offs + sz);
		b = *_v;
		memcpy(b + offs, bytes, sz);

		return sz;
	}
};

void SB4CompressJPEG(const SB4Image *img, unsigned int quality, terVector<unsigned char> &out)
{
	terSafePointer<TDP::terJPEGCompressor> cmp;
	const SB4ImagePlane *planes[3];
	const unsigned char *rows[3];
	unsigned int i, j;
	myVectorWriter vectorWriter(&out);
	terVector<unsigned char> rowBytes;
	unsigned char *rowPtr;

	rowBytes.Alloc(img->Width() * 3);

	for(i=0;i<3;i++)
		planes[i] = img->Plane(i);

	cmp = TDP::NewJPEGCompressor(img->Width(), img->Height(), quality, terJPEGColorSpaceYUV_k, &vectorWriter);

	for(i=0;i<img->Height();i++)
	{
		// Deinterlace rows
		for(j=0;j<3;j++)
			rows[j] = planes[j]->Row(i);
		rowPtr = rowBytes;
		for(j=0;j<img->Width();j++)
		{
			*rowPtr++ = *rows[0]++;
			*rowPtr++ = *rows[1]++;
			*rowPtr++ = *rows[2]++;
		}
		// Add scanline
		cmp->AddScanLine(rowBytes);
	}
	cmp.Discard();
}


class myVectorReader : public TDP::terJPEGReader
{
private:
	const unsigned char *_bytes;
	size_t _nBytes;
public:
	myVectorReader(const unsigned char *bytes, size_t nBytes)
	{
		_bytes = bytes;
		_nBytes = nBytes;
	}

	size_t Read(void *bytes, size_t sz)
	{
		if(_nBytes < sz)
			sz = _nBytes;

		memcpy(bytes, _bytes, sz);
		_nBytes -= sz;
		_bytes += sz;
		return sz;
	}
};


void SB4DecompressJPEG(SB4Image *img, unsigned char *bytes, size_t numBytes)
{
	terSafePointer<TDP::terJPEGDecompressor> cmp;
	myVectorReader reader(bytes, numBytes);
	unsigned int i, j;
	terVector<unsigned char> rowBytes;
	unsigned char *rowPtrs[3];
	const unsigned char *inRowPtr;

	cmp = TDP::NewJPEGDecompressor(TDP::terJPEGColorSpaceYUV_k, &reader);

	img->Init(cmp->Width(), cmp->Height());
	rowBytes.Alloc(cmp->Width() * cmp->Height() * 3);

	for(i=0;i<cmp->Height();i++)
	{
		for(j=0;j<3;j++)
			rowPtrs[j] = img->Plane(j)->Row(i);

		cmp->ReadScanLine(rowBytes);
		inRowPtr = rowBytes;

		for(j=0;j<cmp->Width();j++)
		{
			*rowPtrs[0]++ = *inRowPtr++;
			*rowPtrs[1]++ = *inRowPtr++;
			*rowPtrs[2]++ = *inRowPtr++;
		}
	}
	cmp.Discard();
}

void SB4ProgressiveReduceJPEG(sb4videoenc_t *enc, SB4Image *reconstruct, terVector<unsigned char> &result, bool delta)
{
	unsigned int highQuality = 100;
	unsigned int lowQuality = 0;
	unsigned int midQuality;
	SB4Image resultImage;
	SB4Image deltaImage;
	SB4Image resultDeltaImage;
	terVector<unsigned char> jpegData;
	const unsigned char *row1, *row2, *row3;
	unsigned char *ncrow;
	unsigned int x, y, p, h, w;
	int d;
	__int64 resultDiff;

	h = enc->frame_to_enc->Height();
	w = enc->frame_to_enc->Width();

	if(delta)
	{
		deltaImage.Init(w, h);
		resultDeltaImage.Init(w, h);
		for(y=0;y<h;y++)
		{
			for(p=0;p<3;p++)
			{
				row1 = enc->frame_to_enc->Plane(p)->Row(y);
				row2 = enc->last_frame->Plane(p)->Row(y);
				ncrow = deltaImage.Plane(p)->Row(y);

				for(x=0;x<w;x++)
				{
					d = (((int)(*row1++) - (int)(*row2++)) + 255) / 2;
					if(d < 0) d = 0; else if(d > 255) d = 255;
					*ncrow++ = d;
				}
			}
		}
	}

	while(highQuality - lowQuality > 1)
	{
		midQuality = (highQuality + lowQuality) / 2;
		jpegData.Alloc(8);

		// Compress JPEG
		if(delta)
		{
			SB4CompressJPEG(&deltaImage, midQuality, jpegData);
			SB4DecompressJPEG(&resultDeltaImage, &jpegData[8], jpegData.Count() - 8);
			resultImage.Init(w, h);

			// Un-delta image
			for(y=0;y<h;y++)
			{
				for(p=0;p<3;p++)
				{
					row1 = resultDeltaImage.Plane(p)->Row(y);
					row2 = enc->last_frame->Plane(p)->Row(y);
					ncrow = resultImage.Plane(p)->Row(y);

					for(x=0;x<w;x++)
					{
						d = (int)(*row2++) + ((int)(*row1++) * 2) - 255;
						if(d < 0) d = 0; else if(d > 255) d = 255;
						*ncrow++ = d;
					}
				}
			}
		}
		else
		{
			SB4CompressJPEG(enc->frame_to_enc, midQuality, jpegData);
			SB4DecompressJPEG(&resultImage, &jpegData[8], jpegData.Count() - 8);
		}

		// Get variance
		resultDiff = 0;
		for(y=0;y<h;y++)
		{
			for(p=0;p<3;p++)
			{
				row1 = enc->frame_to_enc->Plane(p)->Row(y);
				row2 = resultImage.Plane(p)->Row(y);

				for(x=0;x<w;x++)
				{
					d = (int)(*row1++) - (int)(*row2++);
					resultDiff += d*d;
				}
			}
		}

		enc->dist = resultDiff;

		if(jpegData.Count() < enc->minBytes)
		{
			// Aim for better quality
			lowQuality = midQuality;
			continue;
		}

		// Are we over budget or not meeting the target dist?
		if(jpegData.Count() > enc->maxBytes || resultDiff < enc->target_dist || jpegData.Count() >= 65536)
		{
			// Aim for reduced quality
			highQuality = midQuality;
			continue;
		}

		// Aim for better quality
		lowQuality = midQuality;
	}

	// Write the header
	ncrow = jpegData;
	if(delta)
		bytestream_put_le16(&ncrow, RoQ_QUAD_DJPEG);
	else
		bytestream_put_le16(&ncrow, RoQ_QUAD_JPEG);
	bytestream_put_le32(&ncrow, jpegData.Count() - 8);
	bytestream_put_le16(&ncrow, 0);


	printf("JPEG P-reduce: Finalized at %i  Dist %f\n", midQuality, (double)enc->dist);

	reconstruct->Init(w, h);
	resultImage.CopyTo(reconstruct);
	jpegData.MoveTo(&result);

	enc->final_size = result.Count();
}
